# -*- coding: utf-8 -*-
import datetime
import os
import sys
import pandas as pd
import numpy as np
from pymongo import MongoClient
from collections import namedtuple
from abc import ABCMeta, abstractmethod
import tushare as ts
import copy

from .error import ParamInitException


class DataHandler(object):
    Bar = namedtuple('Bar', ('symbol', 'datetime', 'open', 'high', 'low', 'close', 'volume'))

    __metaclass__ = ABCMeta

    @abstractmethod
    def get_latest_bars(self, symbol, N=1):
        raise NotImplementedError("Should implement get_latest_bars()!")

    @abstractmethod
    def update_bars(self):
        raise NotImplementedError("Should implement update_bars()!")


class CSVDataHandler(DataHandler):
    def __init__(self, events, stocks, start_date, end_date, params):
        self.events = events

        if 'csv_dir' in params:
            self.csv_dir = params['csv_dir']
        else:
            raise ParamInitException("csv_dir are not found, needed by CSVDataHandler")


        self.symbol_list = stocks
        self.start_date = start_date
        self.end_date = end_date

        self.symbol_data = {}
        self.latest_symbol_data = {}
        self.continue_backtest = True

        self._open_convert_csv_files()

    def _open_convert_csv_files(self):
        comb_index = None
        for s in self.symbol_list:
            self.symbol_data[s] = pd.read_csv(
                os.path.join(self.csv_dir, '%s.csv' % s),
                header=0, index_col=0, parse_dates=True,
                names=['datetime', 'open', 'high', 'low', 'close', 'volume']
            ).sort_index()[self.start_date:self.end_date]
            if comb_index is None:
                comb_index = self.symbol_data[s].index
            else:
                comb_index.union(self.symbol_data[s].index)
            self.latest_symbol_data[s] = []

        for s in self.symbol_list:
            self.symbol_data[s] = self.symbol_data[s].reindex(index=comb_index, method='pad').iterrows()

    def _get_new_bar(self, symbol):
        for b in self.symbol_data[symbol]:
            yield DataHandler.Bar(symbol, b[0], b[1][0], b[1][1], b[1][2], b[1][3], b[1][4])

    def update_bars(self):
        bar_dict = {}
        datetime = None
        for s in self.symbol_list:
            try:
                bar = next(self._get_new_bar(s))
            except StopIteration:
                self.continue_backtest = False
            else:
                if bar is not None:
                    self.latest_symbol_data[s].append(bar)
                    bar_dict[s] = bar
                    datetime = bar.datetime
        self.events.put(BarEvent(bar_dict, datetime))


    def get_latest_bars(self, symbol, N=1):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print(symbol + "isnot avaliable")
        else:
            return bars_list[-N:]

    def get_latest_bar(self, symbol):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print("not avaliable")
        else:
            return bars_list[-1]


class TushareDataHandler(DataHandler):
    def __init__(self, event_hub, stocks, start_date, end_date, params):
        self.events = event_hub
        self.symbol_list = copy.copy(stocks)
        if 'hs300' not in self.symbol_list:
            self.symbol_list.append('hs300')
        # set a reference symbol
        self.reference_symbol = 'hs300'
        self.start_date = start_date
        self.end_date = end_date
        self.now = start_date

        self.symbol_data_cache = {}
        self.symbol_data = {}
        self.is_stop = {}
        self.latest_symbol_data = {}
        self.continue_backtest = True
        self._prepare_reference_data()
        self._prepare_data()

    def _prepare_reference_data(self):
        df = ts.get_hist_data(self.reference_symbol, start=self.start_date.strftime('%Y-%m-%d'), end=self.end_date.strftime('%Y-%m-%d'))
        df['date'] = df.index
        df = df.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume']]
        df = df.sort_index()
        datalist = np.array(df).tolist()
        self.reference_symbol_data = self._generate_data(datalist)
        self.is_stop[self.reference_symbol] = False

    def _prepare_data(self):
        for s in self.symbol_list:
            df = ts.get_hist_data(s, start=self.start_date.strftime('%Y-%m-%d'), end=self.end_date.strftime('%Y-%m-%d'))
            df['date'] = df.index
            df = df.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume']]
            df = df.sort_index()
            datalist = np.array(df).tolist()
            self.symbol_data[s] = self._generate_data(datalist)
            self.latest_symbol_data[s] = []
            self.is_stop[s] = False

    def _prepare_symbol_data(self, s):
        df = ts.get_hist_data(s, start=self.now.strftime('%Y-%m-%d'), end=self.end_date.strftime('%Y-%m-%d'))
        df['date'] = df.index
        df = df.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume']]
        df = df.sort_index()
        datalist = np.array(df).tolist()
        self.symbol_data[s] = self._generate_data(datalist)
        self.latest_symbol_data[s] = []
        self.is_stop[s] = False

    def _generate_data(self, list):
        for each in list:
            yield(each)

    def _get_new_bar(self, symbol):
        for b in self.symbol_data[symbol]:
            yield DataHandler.Bar(symbol, pd.Timestamp(b[0]), b[1], b[2], b[3], b[4], b[5])

    def _get_new_ref(self):
        for b in self.reference_symbol_data:
            yield DataHandler.Bar(self.reference_symbol, pd.Timestamp(b[0]), b[1], b[2], b[3], b[4], b[5])

    def update_ref(self):
        try:
            ref = next(self._get_new_ref())
        except StopIteration:
            self.continue_backtest = False
            return
        self.now = ref.datetime

    def update_bars(self):
        bar_dict = {}

        for s in self.symbol_list:
            if self.is_stop[s]:
                # 没数据了
                if s not in self.symbol_data_cache:
                    continue
                # 把停牌的股票恢复
                if self.symbol_data_cache[s].datetime == self.now:
                    self.is_stop[s] = False
                    bar = copy.copy(self.symbol_data_cache[s])
                    self.latest_symbol_data[s].append(bar)
                    bar_dict[s] = bar
                    del self.symbol_data_cache[s]
                continue

            try:
                bar = next(self._get_new_bar(s))
            except StopIteration:
                # self.continue_backtest = False
                self.is_stop[s] = True  # 可能是停牌了, 也可能是没数据了
            else:
                if bar is not None:
                    if bar.datetime != self.now:
                        self.symbol_data_cache[s] = bar
                        self.is_stop[s] = True
                        continue

                    self.latest_symbol_data[s].append(bar)
                    bar_dict[s] = bar

        if len(bar_dict) == 0:
            self.continue_backtest = False
        if self.continue_backtest:
            return bar_dict

    def update_universe(self, stocks):
        for s in stocks:
            if s not in self.symbol_list:
                self.symbol_list.append(s)
                self._prepare_symbol_data(s)

    def check_stop(self, symbol):
        if self.is_stop[symbol]:
            return True
        return False


    def get_latest_bars(self, symbol, N=1):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print(symbol + "isnot avaliable")
        else:
            return bars_list[-N:]

    def get_latest_bar(self, symbol):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print("not avaliable")
        else:
            return bars_list[-1]


class MongodbDataHandler(DataHandler):
    # 本地数据处理类. 这个相比tushare数据处理类的好处就是在每日更新数据的时候, 可以一条一条更新(通过生成器), 而tushare是需要一次把所有
    # 数据读入到内存中的.
    def __init__(self, event_hub, stocks, start_date, end_date, params):
        # 事件引擎
        self.events = event_hub
        # 当前的股票列表, 只有在这个股票列表里面的股票当天才能够接收到更新, 注意用的是copy, 外面的stocks更新过之后需要执行一次
        # update_universe才能够更新股票.
        # 为什么加入update_universe, 主要是如果用户把股票从stocks删除了, 但是用户并没有卖出股票,这样很可能会造成错误.
        # 所以update_universe只会新增股票, 而不会删除原来的老股票
        self.symbol_list = copy.copy(stocks)
        # 加入沪深300作为参照股票
        # 可能有些股票会存在停牌的情况, 这样的话时间很难处理.
        # 而沪深300肯定不会停牌的, 所以可以把它当成参照时间
        if 'hs300' not in self.symbol_list:
            self.symbol_list.append('hs300')
        # set a reference symbol
        self.reference_symbol = 'hs300'
        # 回测的开始时间和结束时间
        self.start_date = start_date
        self.end_date = end_date
        self.now = start_date

        # 有些股票可能会停牌, 因为数据读取的问题, 读取过得数据不可能再读取了, 所以要找一个地方把它们先存起来
        self.symbol_data_cache = {}
        # 用于数据读取
        self.symbol_data = {}
        # 用于判断是否已经停牌
        self.is_stop = {}
        # 用于获取股票历史的数据
        self.latest_symbol_data = {}
        # 是否可以停止回测
        self.continue_backtest = True
        # mongodb的接口
        self.db = self._get_db()
        # 准备数据
        self._prepare_reference_data()
        self._prepare_data()

    def _get_db(self):
        # 连接mongodb
        client = MongoClient("localhost", 27017)
        db = client.stock
        return db

    def _prepare_reference_data(self):
        # 查询参照股票的数据
        start_date_str = self.start_date.strftime('%Y-%m-%d')
        end_date_str = self.end_date.strftime('%Y-%m-%d')
        self.reference_symbol_data = self.db.stocks.find({'code': self.reference_symbol,
                                                          'date': {'$gte': start_date_str, '$lte': end_date_str}})
        self.is_stop[self.reference_symbol] = False

    def _prepare_data(self):
        # 查询股票数据, 这里数据得到的是一个生成器.
        start_date_str = self.start_date.strftime('%Y-%m-%d')
        end_date_str = self.end_date.strftime('%Y-%m-%d')
        for s in self.symbol_list:
            self.symbol_data[s] = self.db.stocks.find({'code': s,
                                                       'date': {'$gte': start_date_str, '$lte': end_date_str}})
            self.latest_symbol_data[s] = []
            self.is_stop[s] = False

    def _prepare_symbol_data(self, s):
        # 因为update_universe会加入新股票, 加入新股票之后需要准备好新股票的数据
        now_str = self.now.strftime('%Y-%m-%d')
        end_date_str = self.end_date.strftime('%Y-%m-%d')
        self.symbol_data[s] = self.db.stocks.find({'code': s,
                                                   'date': {'$gte': now_str, '$lte': end_date_str}})
        self.latest_symbol_data[s] = []
        self.is_stop[s] = False


    def update_universe(self, stocks):
        # 更新全局的股票列表, 股票数目只会增加,不会减少
        for s in stocks:
            if s not in self.symbol_list:
                self.symbol_list.append(s)
                self._prepare_symbol_data(s)

    def _get_new_bar(self, symbol):
        # 从生成器中获取一条新的数据, 转成我们想要的格式, 并return
        for b in self.symbol_data[symbol]:
            yield DataHandler.Bar(symbol, pd.Timestamp(b['date']),
                                  b['open'], b['high'], b['low'], b['close'], b['volume'])

    def _get_new_ref(self):
        # 获取一条参照股票的数据
        for b in self.reference_symbol_data:
            yield DataHandler.Bar(self.reference_symbol, pd.Timestamp(b['date']),
                                  b['open'], b['high'], b['low'], b['close'], b['volume'])

    def update_ref(self):
        # 更新参照股票的数据,目的主要是为了更新now
        try:
            ref = next(self._get_new_ref())
        except StopIteration:
            self.continue_backtest = False
            return
        self.now = ref.datetime

    def update_bars(self):
        bar_dict = {}
        # 对于每只股票, 读取新数据, 如果新数据和参照股票的时间对不上, 说明该股票那天停牌了, 所以放到cache里面.等下次时间一致之后恢复.
        # 返回当天所有的股票的数据
        for s in self.symbol_list:
            if self.is_stop[s]:
                # 没数据了
                if s not in self.symbol_data_cache:
                    continue
                # 把停牌的股票恢复
                if self.symbol_data_cache[s].datetime == self.now:
                    self.is_stop[s] = False
                    bar = copy.copy(self.symbol_data_cache[s])
                    self.latest_symbol_data[s].append(bar)
                    bar_dict[s] = bar
                    del self.symbol_data_cache[s]
                continue

            try:
                bar = next(self._get_new_bar(s))
            except StopIteration:
                self.is_stop[s] = True  # 可能是停牌了, 也可能是没数据了
            else:
                if bar is not None:
                    if bar.datetime != self.now:
                        self.symbol_data_cache[s] = bar
                        self.is_stop[s] = True
                        continue

                    self.latest_symbol_data[s].append(bar)
                    bar_dict[s] = bar

        if len(bar_dict) == 0:
            self.continue_backtest = False
        if self.continue_backtest:
            return bar_dict

    # 判断该只股票现在是否停牌了
    def check_stop(self, symbol):
        if self.is_stop[symbol]:
            return True
        return False

    # 获取最近N天的股票的数据
    def get_latest_bars(self, symbol, N=1):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print(symbol + "isnot avaliable")
        else:
            return bars_list[-N:]

    # 获取最近一天的股票的数据
    def get_latest_bar(self, symbol):
        try:
            bars_list = self.latest_symbol_data[symbol]
        except KeyError:
            print("not avaliable")
        else:
            return bars_list[-1]